from langchain_core.documents import Document
from langchain_community.vectorstores import Qdrant
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_community.chat_models import ChatZhipuAI
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFium2Loader as PyPDFLoader
from pathlib import Path
import warnings
from typing import List, Dict
import time
warnings.filterwarnings('ignore')

# 配置参数
ZHIPU_API_KEY = "2f39319bdd864fc4a41bf6b8eed6efbc.uIsAkRrMwejVTIyc"
INITIAL_DOC_PATH = '01_“未来校园”智能应用专项赛.pdf'

# 初始化组件
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50
)
embeddings = ZhipuAIEmbeddings(api_key=ZHIPU_API_KEY)

# 新增带有延迟的Embedding类
class DelayedZhipuAIEmbeddings(ZhipuAIEmbeddings):
    """添加API调用延迟的嵌入模型"""
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        embeddings = super().embed_documents(texts)
        time.sleep(2)  # 添加2秒延迟
        return embeddings

    def embed_query(self, text: str) -> List[float]:
        embedding = super().embed_query(text)
        time.sleep(1.5)  # 添加2秒延迟
        return embedding

# 新增带有延迟的Chat模型类
class DelayedChatZhipuAI(ChatZhipuAI):
    """添加API调用延迟的聊天模型"""
    def invoke(self, input, **kwargs):
        result = super().invoke(input, **kwargs)
        time.sleep(2)  # 添加2秒延迟
        return result
embeddings = DelayedZhipuAIEmbeddings(api_key=ZHIPU_API_KEY)  # 使用延迟版本
# 初始化知识库
# 初始化知识库
def initialize_vectorstore() -> Qdrant:
    """初始化带有默认文档的向量存储"""
    loader = PyPDFLoader(INITIAL_DOC_PATH)
    documents = loader.load()
    chunked_docs = text_splitter.split_documents(documents)
    return Qdrant.from_documents(
        documents=chunked_docs,
        embedding=embeddings,
        location=":memory:",
        collection_name="smart_campus_docs"
    )

vectorstore = initialize_vectorstore()


# 增强版文档处理功能
def process_file(file_path: Path) -> Dict:
    """处理单个文件并返回结果"""
    result = {
        'file': str(file_path),
        'status': 'success',
        'chunks': 0,
        'error': None
    }

    try:
        # 文件验证
        if not file_path.exists():
            raise FileNotFoundError(f"文件不存在: {file_path}")
        if file_path.suffix.lower() != '.pdf':
            raise ValueError("仅支持PDF文件格式")

        # 加载和分割文档
        loader = PyPDFLoader(str(file_path))
        docs = loader.load()
        chunks = text_splitter.split_documents(docs)

        # 添加到知识库
        vectorstore.add_documents(chunks)
        result['chunks'] = len(chunks)
        print(f"  ✅ 已加载: {file_path.name} ({len(chunks)}条)")

    except Exception as e:
        result['status'] = 'failed'
        result['error'] = str(e)
        print(f"  ❌ 加载失败: {file_path.name} - {str(e)}")

    return result


def batch_add_documents(paths: List[str]) -> Dict:
    """
    批量添加文档或文件夹
    :param paths: 文件/文件夹路径列表
    :return: 包含处理结果的字典
    """
    total_results = {
        'total_files': 0,
        'success_files': 0,
        'failed_files': 0,
        'total_chunks': 0,
        'errors': []
    }

    processed_files = set()

    for input_path in paths:
        current_path = Path(input_path.strip()).resolve()

        # 处理文件夹
        if current_path.is_dir():
            print(f"\n📂 扫描文件夹: {current_path}")
            for file_path in current_path.glob('**/*.pdf'):
                if file_path in processed_files:
                    continue
                total_results['total_files'] += 1
                res = process_file(file_path)
                processed_files.add(file_path)

                if res['status'] == 'success':
                    total_results['success_files'] += 1
                    total_results['total_chunks'] += res['chunks']
                else:
                    total_results['failed_files'] += 1
                    total_results['errors'].append({
                        'file': res['file'],
                        'error': res['error']
                    })

        # 处理单个文件
        elif current_path.is_file():
            if current_path in processed_files:
                continue
            total_results['total_files'] += 1
            res = process_file(current_path)
            processed_files.add(current_path)

            if res['status'] == 'success':
                total_results['success_files'] += 1
                total_results['total_chunks'] += res['chunks']
            else:
                total_results['failed_files'] += 1
                total_results['errors'].append({
                    'file': res['file'],
                    'error': res['error']
                })

        else:
            error_msg = f"路径不存在: {current_path}"
            print(f"  ❌ {error_msg}")
            total_results['errors'].append({
                'file': str(current_path),
                'error': error_msg
            })

    return total_results


# 构建问答链（修改此处使用延迟版本）
retriever = vectorstore.as_retriever()
qa_chain = (
    RunnableParallel({"context": retriever, "question": RunnablePassthrough()})
    | ChatPromptTemplate.from_template("基于以下上下文回答：\n{context}\n\n问题：{question}")
    | DelayedChatZhipuAI(api_key=ZHIPU_API_KEY, temperature=0.5, model="glm-4")  # 使用延迟版本
    | StrOutputParser()
)

# 交互界面
def interactive_interface():
    """增强版交互界面"""
    print("\n" + "=" * 40)
    print("🏫 智能校园知识助手")
    print("=" * 40)
    print("功能菜单:")
    print("1. 提问\n2. 批量添加文档\n3. 退出系统")

    while True:
        try:
            choice = input("\n请输入选项数字：").strip()

            # 提问模式
            if choice == '1':
                query = input("\n请输入您的问题：").strip()
                if not query:
                    print("⚠️ 问题不能为空")
                    continue
                response = qa_chain.invoke(query)
                print("\n" + "-" * 40)
                print("🤖 系统回答：\n")
                print(response)
                print("\n" + "-" * 40)

            # 批量添加模式
            elif choice == '2':
                print("\n📁 支持以下输入方式：")
                print("- 单个文件路径 (例如: D:/文档/比赛规则.pdf)")
                print("- 多个路径用分号分隔 (例如: path1; path2; path3)")
                print("- 整个文件夹路径 (例如: D:/数据文件夹)")

                input_paths = input("\n请输入文件/文件夹路径：").split(';')
                paths = [p.strip() for p in input_paths if p.strip()]

                if not paths:
                    print("⚠️ 未输入有效路径")
                    continue

                print("\n🚀 开始处理文档...")
                results = batch_add_documents(paths)

                # 显示统计结果
                print("\n" + "=" * 40)
                print("📊 导入统计:")
                print(f"• 发现文件总数: {results['total_files']}")
                print(f"• 成功加载文件: {results['success_files']}")
                print(f"• 失败文件数: {results['failed_files']}")
                print(f"• 新增知识条目: {results['total_chunks']}")

                # 显示错误详情
                if results['errors']:
                    print("\n❌ 错误详情:")
                    for error in results['errors']:
                        print(f"  - 文件: {error['file']}")
                        print(f"    原因: {error['error']}")
                        print("-" * 30)

            # 退出系统
            elif choice == '3':
                print("\n感谢使用，再见！👋")
                break

            else:
                print("⚠️ 无效选项，请重新输入")

        except KeyboardInterrupt:
            print("\n操作已中断")
            break
        except Exception as e:
            print(f"❌ 发生未预期错误: {str(e)}")


if __name__ == "__main__":
    print("🔍 正在初始化知识库...")
    interactive_interface()